"""
CodeQL完整工作流脚本
实现从source到sink的完整代码分析流程
"""

import os
import sys
import subprocess
import json
import csv
from pathlib import Path
from typing import List, Dict, Optional
import argparse


class CodeQLWorkflow:
    """CodeQL完整工作流管理器"""
    
    def __init__(self, codeql_path: str = None):
        self.codeql_path = codeql_path or self._find_codeql()
        self.queries_dir = Path(__file__).parent.parent / "queries"
        self.output_dir = Path(__file__).parent.parent / "output"
        self.output_dir.mkdir(exist_ok=True)
    
    def _find_codeql(self) -> str:
        """查找CodeQL可执行文件"""
        # 常见的CodeQL安装路径
        possible_paths = [
            "codeql",
            "/usr/local/bin/codeql",
            "/opt/homebrew/bin/codeql",
            "C:\\codeql\\codeql.exe"
        ]
        
        for path in possible_paths:
            try:
                result = subprocess.run([path, "version"], 
                                      capture_output=True, text=True, timeout=10)
                if result.returncode == 0:
                    print(f"✅ 找到CodeQL: {path}")
                    return path
            except (subprocess.TimeoutExpired, FileNotFoundError):
                continue
        
        raise FileNotFoundError("未找到CodeQL可执行文件，请确保已安装并配置PATH")
    
    def create_database(self, 
                       source_dir: str, 
                       language: str = "cpp",
                       output_db: str = None) -> str:
        """
        创建CodeQL数据库
        
        Args:
            source_dir: 源代码目录
            language: 编程语言
            output_db: 输出数据库路径
            
        Returns:
            数据库路径
        """
        if not output_db:
            output_db = self.output_dir / f"{language}_database"
        
        print(f"🔨 创建CodeQL数据库...")
        print(f"   源代码: {source_dir}")
        print(f"   语言: {language}")
        print(f"   输出: {output_db}")
        
        cmd = [
            self.codeql_path, "database", "create",
            str(output_db),
            f"--language={language}",
            f"--source-root={source_dir}",
            "--command=make"  # 假设使用make构建
        ]
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
            if result.returncode == 0:
                print(f"✅ 数据库创建成功: {output_db}")
                return str(output_db)
            else:
                print(f"❌ 数据库创建失败: {result.stderr}")
                return None
        except subprocess.TimeoutExpired:
            print("❌ 数据库创建超时")
            return None
    
    def run_query(self, 
                  database: str, 
                  query_file: str, 
                  output_file: str = None,
                  format: str = "csv") -> str:
        """
        运行CodeQL查询
        
        Args:
            database: 数据库路径
            query_file: 查询文件路径
            output_file: 输出文件路径
            format: 输出格式 (csv, sarif, json)
            
        Returns:
            输出文件路径
        """
        if not output_file:
            query_name = Path(query_file).stem
            output_file = self.output_dir / f"{query_name}.{format}"
        
        print(f"🔍 运行查询: {query_file}")
        print(f"   数据库: {database}")
        print(f"   输出: {output_file}")
        
        cmd = [
            self.codeql_path, "database", "analyze",
            database,
            query_file,
            f"--format={format}",
            f"--output={output_file}"
        ]
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
            if result.returncode == 0:
                print(f"✅ 查询执行成功: {output_file}")
                return str(output_file)
            else:
                print(f"❌ 查询执行失败: {result.stderr}")
                return None
        except subprocess.TimeoutExpired:
            print("❌ 查询执行超时")
            return None
    
    def run_workflow(self, 
                    source_dir: str,
                    language: str = "cpp",
                    queries: List[str] = None) -> Dict[str, str]:
        """
        运行完整的CodeQL工作流
        
        Args:
            source_dir: 源代码目录
            language: 编程语言
            queries: 要运行的查询列表
            
        Returns:
            查询结果文件路径字典
        """
        if queries is None:
            queries = [
                "CallGraph.ql",
                "DataFlowAnalysis.ql", 
                "ControlFlowAnalysis.ql",
                "SecurityVulnerabilities.ql"
            ]
        
        print("🚀 开始CodeQL完整工作流")
        print("=" * 60)
        
        # 1. 创建数据库
        database = self.create_database(source_dir, language)
        if not database:
            return {}
        
        # 2. 运行查询
        results = {}
        for query in queries:
            query_path = self.queries_dir / query
            if query_path.exists():
                output_file = self.run_query(database, str(query_path))
                if output_file:
                    results[query] = output_file
            else:
                print(f"⚠️  查询文件不存在: {query_path}")
        
        print(f"\n✅ 工作流完成，生成了 {len(results)} 个结果文件")
        return results
    
    def analyze_results(self, results: Dict[str, str]) -> Dict[str, any]:
        """
        分析查询结果
        
        Args:
            results: 查询结果文件路径字典
            
        Returns:
            分析结果
        """
        analysis = {
            "call_graph": {},
            "data_flow": {},
            "control_flow": {},
            "security_vulnerabilities": {}
        }
        
        for query, output_file in results.items():
            if not os.path.exists(output_file):
                continue
            
            print(f"📊 分析结果: {query}")
            
            try:
                with open(output_file, 'r', encoding='utf-8') as f:
                    if output_file.endswith('.csv'):
                        reader = csv.DictReader(f)
                        rows = list(reader)
                        
                        if "CallGraph" in query:
                            analysis["call_graph"] = self._analyze_call_graph(rows)
                        elif "DataFlow" in query:
                            analysis["data_flow"] = self._analyze_data_flow(rows)
                        elif "ControlFlow" in query:
                            analysis["control_flow"] = self._analyze_control_flow(rows)
                        elif "Security" in query:
                            analysis["security_vulnerabilities"] = self._analyze_security(rows)
                        
                        print(f"   📈 发现 {len(rows)} 条记录")
            
            except Exception as e:
                print(f"   ❌ 分析失败: {e}")
        
        return analysis
    
    def _analyze_call_graph(self, rows: List[Dict]) -> Dict:
        """分析调用图结果"""
        if not rows:
            return {}
        
        # 统计调用关系
        call_count = len(rows)
        unique_functions = set()
        unique_files = set()
        
        for row in rows:
            unique_functions.add(row.get('caller_qualified', ''))
            unique_functions.add(row.get('callee_qualified', ''))
            unique_files.add(row.get('caller_file', ''))
            unique_files.add(row.get('callee_file', ''))
        
        return {
            "total_calls": call_count,
            "unique_functions": len(unique_functions),
            "unique_files": len(unique_files),
            "functions": list(unique_functions)[:10]  # 前10个函数
        }
    
    def _analyze_data_flow(self, rows: List[Dict]) -> Dict:
        """分析数据流结果"""
        if not rows:
            return {}
        
        # 统计数据流类型
        flow_types = {}
        for row in rows:
            flow_type = row.get('flow_path', 'unknown')
            flow_types[flow_type] = flow_types.get(flow_type, 0) + 1
        
        return {
            "total_flows": len(rows),
            "flow_types": flow_types,
            "sources": len(set(row.get('source_location', '') for row in rows)),
            "sinks": len(set(row.get('sink_location', '') for row in rows))
        }
    
    def _analyze_control_flow(self, rows: List[Dict]) -> Dict:
        """分析控制流结果"""
        if not rows:
            return {}
        
        # 统计控制流类型
        flow_types = {}
        for row in rows:
            flow_type = row.get('control_flow_type', 'unknown')
            flow_types[flow_type] = flow_types.get(flow_type, 0) + 1
        
        return {
            "total_flows": len(rows),
            "flow_types": flow_types,
            "functions": len(set(row.get('function_name', '') for row in rows))
        }
    
    def _analyze_security(self, rows: List[Dict]) -> Dict:
        """分析安全漏洞结果"""
        if not rows:
            return {}
        
        # 统计漏洞类型和严重程度
        vulnerability_types = {}
        severity_levels = {}
        
        for row in rows:
            vuln_type = row.get('vulnerability_type', 'unknown')
            severity = row.get('severity', 'unknown')
            
            vulnerability_types[vuln_type] = vulnerability_types.get(vuln_type, 0) + 1
            severity_levels[severity] = severity_levels.get(severity, 0) + 1
        
        return {
            "total_vulnerabilities": len(rows),
            "vulnerability_types": vulnerability_types,
            "severity_levels": severity_levels,
            "critical_count": severity_levels.get('Critical', 0),
            "high_count": severity_levels.get('High', 0),
            "medium_count": severity_levels.get('Medium', 0)
        }
    
    def generate_report(self, analysis: Dict[str, any], output_file: str = None) -> str:
        """
        生成分析报告
        
        Args:
            analysis: 分析结果
            output_file: 输出文件路径
            
        Returns:
            报告文件路径
        """
        if not output_file:
            output_file = self.output_dir / "codeql_analysis_report.md"
        
        print(f"📝 生成分析报告: {output_file}")
        
        report_content = self._create_report_content(analysis)
        
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(report_content)
        
        print(f"✅ 报告生成成功: {output_file}")
        return str(output_file)
    
    def _create_report_content(self, analysis: Dict[str, any]) -> str:
        """创建报告内容"""
        content = """# CodeQL 代码分析报告

## 📊 分析概览

本报告基于CodeQL的source到sink分析模式，对代码进行了全面的静态分析。

"""
        
        # 调用图分析
        if analysis.get("call_graph"):
            cg = analysis["call_graph"]
            content += f"""## 🔗 调用图分析

- **总调用关系数**: {cg.get('total_calls', 0)}
- **唯一函数数**: {cg.get('unique_functions', 0)}
- **涉及文件数**: {cg.get('unique_files', 0)}

### 主要函数
"""
            for func in cg.get('functions', [])[:5]:
                content += f"- {func}\n"
            content += "\n"
        
        # 数据流分析
        if analysis.get("data_flow"):
            df = analysis["data_flow"]
            content += f"""## 🌊 数据流分析

- **总数据流数**: {df.get('total_flows', 0)}
- **数据源数量**: {df.get('sources', 0)}
- **数据汇数量**: {df.get('sinks', 0)}

### 数据流类型分布
"""
            for flow_type, count in df.get('flow_types', {}).items():
                content += f"- {flow_type}: {count}\n"
            content += "\n"
        
        # 控制流分析
        if analysis.get("control_flow"):
            cf = analysis["control_flow"]
            content += f"""## 🔄 控制流分析

- **总控制流数**: {cf.get('total_flows', 0)}
- **涉及函数数**: {cf.get('functions', 0)}

### 控制流类型分布
"""
            for flow_type, count in cf.get('flow_types', {}).items():
                content += f"- {flow_type}: {count}\n"
            content += "\n"
        
        # 安全漏洞分析
        if analysis.get("security_vulnerabilities"):
            sv = analysis["security_vulnerabilities"]
            content += f"""## 🔒 安全漏洞分析

- **总漏洞数**: {sv.get('total_vulnerabilities', 0)}
- **严重漏洞**: {sv.get('critical_count', 0)}
- **高危漏洞**: {sv.get('high_count', 0)}
- **中危漏洞**: {sv.get('medium_count', 0)}

### 漏洞类型分布
"""
            for vuln_type, count in sv.get('vulnerability_types', {}).items():
                content += f"- {vuln_type}: {count}\n"
            content += "\n"
        
        content += """## 🎯 建议

1. **代码质量**: 基于调用图分析，建议重构复杂度过高的函数
2. **数据安全**: 关注数据流分析中发现的不安全数据传递
3. **控制流**: 优化控制流结构，提高代码可读性
4. **安全加固**: 优先修复严重和高危安全漏洞

---
*报告生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
"""
        
        return content


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="CodeQL完整工作流")
    parser.add_argument("--source", required=True, help="源代码目录")
    parser.add_argument("--language", default="cpp", help="编程语言")
    parser.add_argument("--codeql-path", help="CodeQL可执行文件路径")
    parser.add_argument("--queries", nargs="+", help="要运行的查询文件")
    
    args = parser.parse_args()
    
    try:
        # 创建工作流管理器
        workflow = CodeQLWorkflow(args.codeql_path)
        
        # 运行完整工作流
        results = workflow.run_workflow(
            source_dir=args.source,
            language=args.language,
            queries=args.queries
        )
        
        if results:
            # 分析结果
            analysis = workflow.analyze_results(results)
            
            # 生成报告
            report_file = workflow.generate_report(analysis)
            
            print(f"\n🎉 工作流完成！")
            print(f"📊 分析结果: {len(results)} 个查询")
            print(f"📝 报告文件: {report_file}")
        else:
            print("❌ 工作流执行失败")
    
    except Exception as e:
        print(f"❌ 工作流执行出错: {e}")


if __name__ == "__main__":
    main()
